import pandas as pd
import numpy as np
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
import json
from typing import Dict, Any, Tuple
import socket
import pickle
import logging

logger = logging.getLogger(__name__)

class DataPreparator:
    def __init__(self, fernet_key: bytes):
        self.fernet_key = fernet_key
        self.cipher_suite = Fernet(self.fernet_key)
        
    def load_and_prepare_data(self, file_path: str) -> pd.DataFrame:
        """加载并准备数据"""
        # 根据文件类型加载数据
        if file_path.endswith('.csv'):
            data = pd.read_csv(file_path)
            # 转换日期列
            if 'date' in data.columns:
                data['date'] = pd.to_datetime(data['date'])
            # 确保数值列的类型正确
            numeric_columns = ['temperature', 'humidity', 'wind_speed', 'precipitation', 
                             'pressure', 'uv_index', 'visibility', 'cloud_cover', 
                             'air_quality_index']
            for col in numeric_columns:
                if col in data.columns:
                    data[col] = pd.to_numeric(data[col], errors='coerce')
        elif file_path.endswith('.json'):
            data = pd.read_json(file_path)
        else:
            raise ValueError("不支持的文件格式")
            
        return data
    
    def encrypt_data(self, data: pd.DataFrame) -> bytes:
        """加密数据"""
        data_bytes = data.to_json().encode()
        return self.cipher_suite.encrypt(data_bytes)
    
    def extract_metadata(self, data: pd.DataFrame) -> Dict[str, Any]:
        """提取元数据"""
        metadata = {
            "columns": list(data.columns),
            "dtypes": {col: str(data[col].dtype) for col in data.columns},
            "row_count": len(data),
            "null_counts": data.isnull().sum().to_dict(),
            "column_stats": {}
        }
        
        for col in data.columns:
            if pd.api.types.is_numeric_dtype(data[col]):
                metadata["column_stats"][col] = {
                    "min": float(data[col].min()),
                    "max": float(data[col].max()),
                    "mean": float(data[col].mean())
                }
            elif pd.api.types.is_datetime64_any_dtype(data[col]):
                metadata["column_stats"][col] = {
                    "start_date": str(data[col].min()),
                    "end_date": str(data[col].max())
                }
        
        return metadata
    
    def hash_metadata(self, metadata: Dict[str, Any]) -> bytes:
        """生成元数据哈希"""
        metadata_str = json.dumps(metadata, sort_keys=True)
        digest = hashes.Hash(hashes.SHA256())
        digest.update(metadata_str.encode())
        return digest.finalize()
    
    def verify_format_compliance(self, data: pd.DataFrame, 
                               format_rules: Dict[str, str]) -> Dict[str, float]:
        """验证格式合规性"""
        compliance = {}
        for col, rule in format_rules.items():
            if col not in data.columns:
                logger.warning(f"数据中不存在规则指定的列: {col}")
                continue
                
            if rule == "date":
                try:
                    pd.to_datetime(data[col])
                    compliance[col] = 1.0
                except Exception as e:
                    logger.warning(f"列 {col} 的日期格式验证失败: {str(e)}")
                    compliance[col] = 0.0
            elif rule == "numeric":
                try:
                    pd.to_numeric(data[col])
                    # 添加范围检查
                    if col == 'temperature':
                        compliance[col] = ((data[col] >= -50) & (data[col] <= 50)).mean()
                    elif col == 'humidity':
                        compliance[col] = ((data[col] >= 0) & (data[col] <= 100)).mean()
                    elif col == 'wind_speed':
                        compliance[col] = (data[col] >= 0).mean()
                    elif col == 'precipitation':
                        compliance[col] = (data[col] >= 0).mean()
                    elif col == 'pressure':
                        compliance[col] = ((data[col] >= 900) & (data[col] <= 1100)).mean()
                    elif col == 'uv_index':
                        compliance[col] = ((data[col] >= 0) & (data[col] <= 15)).mean()
                    elif col == 'visibility':
                        compliance[col] = (data[col] >= 0).mean()
                    elif col == 'cloud_cover':
                        compliance[col] = ((data[col] >= 0) & (data[col] <= 100)).mean()
                    elif col == 'air_quality_index':
                        compliance[col] = (data[col] >= 0).mean()
                    else:
                        compliance[col] = 1.0
                except Exception as e:
                    logger.warning(f"列 {col} 的数值格式验证失败: {str(e)}")
                    compliance[col] = 0.0
            else:
                logger.warning(f"未知的格式规则: {rule}")
                compliance[col] = 0.0
        return compliance
    
    def encrypt_compliance_results(self, compliance: Dict[str, float]) -> bytes:
        """加密合规性结果"""
        try:
            # 确保所有值都是浮点数
            compliance_float = {k: float(v) for k, v in compliance.items()}
            compliance_bytes = json.dumps(compliance_float).encode()
            return self.cipher_suite.encrypt(compliance_bytes)
        except Exception as e:
            logger.error(f"加密合规性结果时出错: {str(e)}")
            raise
    
    def send_to_central_server(self, host: str, port: int, 
                             encrypted_data: bytes, 
                             metadata_hash: bytes,
                             encrypted_compliance: bytes) -> None:
        """发送数据到中央服务器"""
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.connect((host, port))
                data_to_send = {
                    'encrypted_data': encrypted_data,
                    'metadata_hash': metadata_hash,
                    'encrypted_compliance': encrypted_compliance
                }
                # 添加数据验证
                if not all(key in data_to_send for key in ['encrypted_data', 'metadata_hash', 'encrypted_compliance']):
                    raise ValueError("发送的数据缺少必要字段")
                    
                s.sendall(pickle.dumps(data_to_send))
                logger.info("数据成功发送到中央服务器")
        except Exception as e:
            logger.error(f"发送数据到中央服务器时出错: {str(e)}")
            raise 